利益揭露:我在做 UsageRail,一個幫 API 做 key、限流和用量計費的服務,最後一節會提到它。前面的程式都只用 Django 和 DRF,在 Django 6.1、DRF 3.18 上跑過。
API 開始收費以後,「這個 key 一分鐘能打幾次」和「這個客戶這個月還剩多少」會被放在同一張需求單上。兩件事長得很像,都是數請求、超過就回 429,但實際上差很多:
這兩個東西該存的地方不一樣,該在什麼時候記一筆也不一樣。DRF 的 throttling 比較適合前者,硬拿來做後者會踩到幾個坑。下面從 DRF 內建的機制講起,再一步步補上按客戶的限流和額度。
DRF 的 throttle 是在 APIView.initial() 裡跑的,順序是認證、權限、限流,全過了才進到你的 view。內建的 AnonRateThrottle、UserRateThrottle、ScopedRateThrottle 都繼承 SimpleRateThrottle,做法是在 Django cache 裡為每個識別碼存一串時間戳:讀出來、丟掉超過窗口的、數數看還剩幾格,沒超過就把這次的時間戳加上去,再寫回 cache。
拿來做「按客戶」的限流,有幾個地方不夠用:
DEFAULT_THROTTLE_RATES 是 {"user": "100/min"} 這種格式,一個 scope 一個數字。免費方案每分鐘 60 次、付費方案 600 次,沒辦法直接表達。SimpleRateThrottle.__init__() 會在建構時就讀設定、解析速率,那時候還不知道請求是誰發的。CACHES 時用的是 LocMemCache,gunicorn 開 4 個 worker,每個 worker 各數各的,限制等於變成 4 倍。check_throttles() 長這樣:def check_throttles(self, request):
throttle_durations = []
for throttle in self.get_throttles():
if not throttle.allow_request(request, self):
throttle_durations.append(throttle.wait())
...
它不會在第一個失敗的 throttle 就停下,而是每個都問一遍,最後取最長的等待時間。如果你把「扣額度」寫成一個 throttle,跟限流 throttle 並排,那麼限流擋下來的請求,額度一樣會被扣掉。
所以接下來的做法是:限流留在 throttle 裡,換成按客戶、存在共用 cache;額度從 throttle 拿出來,放到資料庫,在限流通過之後才扣。
限流和額度都需要知道「這是哪個客戶」。假設已經有這幾張表:
# billing/models.py
from django.db import models
class Plan(models.Model):
name = models.CharField(max_length=50)
# 留空代表不限
requests_per_minute = models.PositiveIntegerField(null=True, blank=True)
monthly_quota = models.PositiveIntegerField(null=True, blank=True)
class Customer(models.Model):
name = models.CharField(max_length=100)
plan = models.ForeignKey(Plan, on_delete=models.PROTECT)
class ApiKey(models.Model):
customer = models.ForeignKey(Customer, on_delete=models.CASCADE, related_name="api_keys")
key_id = models.CharField(max_length=16, unique=True)
secret_hash = models.CharField(max_length=64)
revoked_at = models.DateTimeField(null=True, blank=True)
認證類別把 key 對應到 ApiKey,放進 request.auth:
# billing/authentication.py
import hashlib
import hmac
from rest_framework import authentication, exceptions, permissions
from .models import ApiKey
class ApiKeyAuthentication(authentication.BaseAuthentication):
def authenticate(self, request):
header = request.headers.get("Authorization", "")
if not header.startswith("Bearer "):
return None
# 假設 key 長這樣:{key_id}.{secret}
key_id, _, secret = header.removeprefix("Bearer ").partition(".")
api_key = (
ApiKey.objects.select_related("customer__plan")
.filter(key_id=key_id, revoked_at__isnull=True)
.first()
)
secret_hash = hashlib.sha256(secret.encode()).hexdigest()
if api_key is None or not hmac.compare_digest(api_key.secret_hash, secret_hash):
raise exceptions.AuthenticationFailed("invalid_api_key")
# 呼叫方是「客戶」,不是 Django 的 User,所以 user 留 None,客戶資料放在 request.auth
return (None, api_key)
def authenticate_header(self, request):
# 有這行,沒帶 key 時 DRF 回 401 而不是 403
return "Bearer"
class HasApiKey(permissions.BasePermission):
def has_permission(self, request, view):
return request.auth is not None
select_related("customer__plan") 不是順手加的。後面的 throttle 和額度都要讀 request.auth.customer.plan,沒有它,每個請求會多兩次查詢。
因為 user 是 None,內建的 IsAuthenticated 會一律拒絕,所以另外寫了 HasApiKey。
改動最少的方式,是讓 throttle 等拿到 request 之後才決定速率。DRF 自己的 ScopedRateThrottle 就是這樣做的:覆寫 __init__() 不讀設定,在 allow_request() 裡才設定 rate。
# billing/throttling.py
from rest_framework.throttling import SimpleRateThrottle
class PlanRateThrottle(SimpleRateThrottle):
scope = "plan"
def __init__(self):
# 不在建構時讀 DEFAULT_THROTTLE_RATES,等拿到 request 才知道是哪個方案
pass
def allow_request(self, request, view):
limit = request.auth.customer.plan.requests_per_minute if request.auth else None
self.rate = f"{limit}/min" if limit else None
self.num_requests, self.duration = self.parse_rate(self.rate)
return super().allow_request(request, view)
def get_cache_key(self, request, view):
return self.cache_format % {"scope": self.scope, "ident": request.auth.customer_id}
識別碼用的是 customer_id,不是 key。一個客戶發了三把 key,三把加起來才是他的速率。
throttle 物件是每個請求新建一個的(get_throttles() 每次都 throttle()),所以在 allow_request() 裡改 self.rate 不會影響別的請求。
這個做法的限制就是第一節講的第 2、3 點:時間戳串的讀寫不是原子的,而且每個客戶在 cache 裡存一串長度等於上限的 list,上限 600 就是 600 個浮點數。
如果要原子性,最簡單的是固定窗口計數:每分鐘一個 key,INCR 一下看有沒有超過。
# billing/throttling.py
import time
from django.core.cache import cache
from rest_framework.throttling import BaseThrottle
class CustomerRateThrottle(BaseThrottle):
window_seconds = 60
def allow_request(self, request, view):
if request.auth is None:
return True
limit = request.auth.customer.plan.requests_per_minute
if limit is None:
return True
now = time.time()
window_start = int(now // self.window_seconds) * self.window_seconds
cache_key = f"ratelimit:{request.auth.customer_id}:{window_start}"
# add 只在 key 不存在時寫入;接著 incr 在 Redis 上是原子操作
cache.add(cache_key, 0, timeout=self.window_seconds * 2)
try:
count = cache.incr(cache_key)
except ValueError:
# key 剛好在 add 和 incr 之間過期
cache.set(cache_key, 1, timeout=self.window_seconds * 2)
count = 1
self.retry_after = window_start + self.window_seconds - now
return count <= limit
def wait(self):
return self.retry_after
設定裡把 cache 換成 Redis(Django 4.0 起內建 RedisCache,不用另外裝 django-redis):
# settings.py
CACHES = {
"default": {
"BACKEND": "django.core.cache.backends.redis.RedisCache",
"LOCATION": "redis://127.0.0.1:6379",
}
}
REST_FRAMEWORK = {
"DEFAULT_AUTHENTICATION_CLASSES": ["billing.authentication.ApiKeyAuthentication"],
"DEFAULT_PERMISSION_CLASSES": ["billing.authentication.HasApiKey"],
"DEFAULT_THROTTLE_CLASSES": ["billing.throttling.CustomerRateThrottle"],
"EXCEPTION_HANDLER": "billing.handlers.api_exception_handler",
}
固定窗口的代價是邊界:上限 60 次的客戶,可以在 12:00:59 打 60 次、12:01:00 再打 60 次,兩秒內 120 次。如果你在意這個,可以改成 token bucket,用一段 Lua script 在 Redis 裡一次算完「補多少 token、扣一個、剩多少」;或者再加一個每秒的上限把突發壓住。只是要防止某個客戶的程式寫壞了狂打的話,固定窗口通常就夠了。
另外,這裡被擋下的請求也會讓計數加一。窗口一到就歸零,所以不會累積;但這代表一個持續狂打的客戶,在整個窗口內都會被擋著,這通常正是你要的。
額度不要放在 cache。理由有三個:
maxmemory-policy allkeys-lru 之類的淘汰策略時,冷門客戶的計數可能被清掉,等於免費送他一個月。所以額度用一張表記,一個客戶一個帳期一列:
# billing/models.py
class UsageCounter(models.Model):
customer = models.ForeignKey(Customer, on_delete=models.CASCADE)
period_start = models.DateField()
used = models.PositiveBigIntegerField(default=0)
class Meta:
constraints = [
models.UniqueConstraint(fields=["customer", "period_start"], name="one_counter_per_period"),
]
扣額度的關鍵是把「還有沒有剩」和「加一」放在同一條 UPDATE:
# billing/quota.py
import datetime as dt
from django.db.models import F
from django.utils import timezone
from rest_framework.exceptions import Throttled
from .models import UsageCounter
class QuotaExceeded(Throttled):
default_detail = "Monthly quota exceeded."
default_code = "quota_exceeded"
def current_period():
# 用 settings.TIME_ZONE 的日曆月,例如 Asia/Taipei
start = timezone.localdate().replace(day=1)
next_start = (start + dt.timedelta(days=32)).replace(day=1)
return start, next_start
def reserve_unit(customer):
start, next_start = current_period()
counter, _ = UsageCounter.objects.get_or_create(customer=customer, period_start=start)
rows = UsageCounter.objects.filter(pk=counter.pk)
quota = customer.plan.monthly_quota
if quota is not None:
rows = rows.filter(used__lt=quota)
# 條件和加一在同一條 UPDATE 裡,兩個請求同時進來也不會一起擠過最後一格
if rows.update(used=F("used") + 1) == 0:
reset_at = timezone.make_aware(dt.datetime.combine(next_start, dt.time.min))
raise QuotaExceeded(wait=(reset_at - timezone.now()).total_seconds())
return counter.pk
def release_unit(counter_pk):
UsageCounter.objects.filter(pk=counter_pk, used__gt=0).update(used=F("used") - 1)
產生的 SQL 大概是 UPDATE ... SET used = used + 1 WHERE id = %s AND used < %s,回傳受影響的列數。0 列就是額度用完了。不用 select_for_update() 先鎖再改,也不用在 Python 裡比大小。
帳期的時區要先想清楚。timezone.localdate() 用的是 settings.TIME_ZONE,設成 Asia/Taipei 的話,每月 1 號台灣時間午夜重置。如果客戶是按 Stripe 訂閱週期付費,帳期應該跟著訂閱走(例如每月 17 號),這時 period_start 要從訂閱資料算,不能用日曆月。
這個設計的弱點是熱點列。一個每秒幾百個請求的大客戶,所有請求都在 UPDATE 同一列,PostgreSQL 的列鎖會讓它們排隊。量大到這個程度時,常見的做法是批次預扣:一次從資料庫扣 100 個單位,放在 Redis 或行程記憶體裡慢慢用,用完再扣下一批,結束時把沒用到的還回去。代價是要處理「行程掛掉時手上那批怎麼辦」。量沒到那裡之前,一條條件 UPDATE 就夠了。
這是整件事最需要想清楚的地方。大致有三種時機:
| 時機 | 優點 | 問題 |
|---|---|---|
| 檢查時就記(DRF throttle 的做法) | 簡單,超過就擋 | 400、500 也被扣 |
| 回應之後才記 | 只算成功的 | 並發時會超額:10 個請求同時檢查都看到「還剩 1」,全部放行 |
| 先預扣,結束時結算 | 不會超額,失敗的會還回去 | 要多寫一段結算 |
限流用第一種沒問題。額度和計費用第三種:進 view 之前先扣一格,回應是 4xx、5xx 就還回去。
在 DRF 裡,initial() 和 finalize_response() 剛好是這兩個掛點:
# billing/quota.py
class QuotaMixin:
def initial(self, request, *args, **kwargs):
# 先跑完認證、權限、限流,全部通過才扣額度
super().initial(request, *args, **kwargs)
self.quota_counter = reserve_unit(request.auth.customer)
def finalize_response(self, request, response, *args, **kwargs):
response = super().finalize_response(request, response, *args, **kwargs)
counter = getattr(self, "quota_counter", None)
if counter is not None and response.status_code >= 400:
release_unit(counter)
return response
用法:
from rest_framework.response import Response
from rest_framework.views import APIView
from billing.quota import QuotaMixin
class Generate(QuotaMixin, APIView):
def post(self, request):
...
return Response({"ok": True})
幾個細節:
super().initial() 跑完才扣,所以被限流擋下的請求不會扣到額度,第一節第 5 點的問題就沒了。reserve_unit() 直接丟 QuotaExceeded,self.quota_counter 沒被設定,finalize_response() 不會誤還。pk,所以 1 月 31 日 23:59:59 進來、2 月 1 日才結束的請求,會還到 1 月那一列,不會還錯月份。finalize_response() 只在 DRF 有處理到的例外路徑上會被呼叫。view 裡丟出 DRF 不認得的例外(例如一個普通的 KeyError)時,handle_exception() 會把它往外丟,finalize_response() 不會跑到,這一格就不會還。如果你很在意 500 不扣額度,可以覆寫 dispatch(),用 try/except 包起來自己還。DRF 對 Throttled 例外有內建處理:狀態碼 429,wait 不是 None 就會加上 Retry-After 標頭(無條件進位成整數秒)。限流 throttle 的 wait() 回傳值、QuotaExceeded(wait=...) 傳進去的秒數,都會走這條路。
預設的 body 是 {"detail": "Request was throttled. Expected available in 54 seconds."},給人看可以,給程式判斷不太方便。用一個例外處理器把錯誤代碼放出來:
# billing/handlers.py
from rest_framework.exceptions import Throttled
from rest_framework.views import exception_handler
def api_exception_handler(exc, context):
response = exception_handler(exc, context)
if response is not None and isinstance(exc, Throttled):
response.data = {"error": exc.get_codes(), "retry_after": exc.wait}
return response
get_codes() 會回傳例外的 default_code,限流是 throttled,額度是 quota_exceeded。
這個區分很重要。限流的 Retry-After 是幾秒,額度的 Retry-After 是「到下個帳期還有幾秒」,月中被擋的話是一百多萬秒。用戶端如果寫了「收到 429 就照 Retry-After 睡一下再重試」,遇到後者會直接卡住好幾天。body 裡有 quota_exceeded,用戶端才知道這不是等一下就好,要升級方案或等下個月。
也有人額度用完回 402 或 403。402 比較適合預付點數用完(要付錢),403 容易跟權限錯誤混在一起。我自己傾向額度用完也回 429,用 body 的錯誤代碼區分。
用上面的程式,建一個方案:每分鐘 5 次、每月 3 次。連打 7 次 POST /v1/generate,第一次故意送壞掉的 body:
# 狀態 Retry-After body 已用額度
0 400 - {"error": "bad_input"} 0
1 200 - {"ok": true} 1
2 200 - {"ok": true} 2
3 200 - {"ok": true} 3
4 429 1386054 {"error": "quota_exceeded", "retry_after": 1386054} 3
5 429 54 {"error": "throttled", "retry_after": 54} 3
6 429 54 {"error": "throttled", "retry_after": 54} 3
可以看到:
Retry-After 是到下個月 1 號台灣時間午夜的秒數。上面的 key、限流、額度和用量計數都得自己維護,做法 B 還要一台 Redis,月底送帳和對帳也得另外寫。UsageRail 把這些拆成兩半:Django 中介軟體負責每個請求的檢查和計數,在 view 之前檢查 key、限流和額度,回應關閉之後才記用量,2xx、3xx 算計費、4xx 和 5xx 不算;送進 Stripe 和事後對帳則由 UsageRail 的服務端做。限流是每個客戶一個 token bucket,存在各行程的記憶體裡;額度則是由服務端把客戶剩下的額度切一小份給每台伺服器,伺服器在自己那份裡扣,大部分請求不用連網路,也就沒有第四節的熱點列問題。用量每 60 秒送一次,帳期結束後記到你自己的 Stripe 帳號開出的發票上。
# settings.py
import os
MIDDLEWARE = [
# ...
"django.contrib.auth.middleware.AuthenticationMiddleware",
"usagerail.django.UsageRailMiddleware",
]
USAGERAIL = {
"app_key": os.environ["USAGERAIL_APP_KEY"],
"auth": "api-key",
}
跟 DRF 一起用時有兩點要注意:它在 DRF 之前就把請求擋下,所以 DRF 的例外處理器看不到這些 429,要改回應格式得用它的 on_reject;另外它讀的是 Authorization: Bearer 或 x-api-key,不讀 DRF TokenAuthentication 那種 Authorization: Token。設定方式、router 產生的 URL 怎麼記錄、行程被強制結束時還沒送出的用量會怎樣,都寫在 Django API 的計量與限流設定裡。